import numpy as np
import ikfastpy
import grpc
import IkfastServer_pb2
import IkfastServer_pb2_grpc
from concurrent import futures
import logging

# Initialize kinematics for UR5 robot arm
ur5_kin = ikfastpy.PyKinematics()
n_joints = ur5_kin.getDOF()


class IkFastServer(IkfastServer_pb2_grpc.Ikfast):

    def ComputeIk(self, request, context):
        result = self.computeIkFast(request)
        return IkfastServer_pb2.ResponseBody(result=result, size=1)

    def ComputeIkList(self, request, context):
        result = self.computeIkFast(request)
        return IkfastServer_pb2.ResponseBody(result=result, size=1)

    def computeIkFast(self, request):
        request_array = request.poseList
        result_list = []
        for posture in request_array:
            matrix = posture.matrix
            request_id = posture.id
            ee = []
            for ma in matrix:
                ee.append(ma)
            result_k = ur5_kin.inverse(ee)
            n_solutions = int(len(result_k) / n_joints)
            joint_configs = np.asarray(result_k).reshape(n_solutions, n_joints)
            joint_result_list = []
            for result in joint_configs:
                joint_result = IkfastServer_pb2.JointsAngle(ax=result)
                joint_result_list.append(joint_result)
            result = IkfastServer_pb2.InverseResult(angles=joint_result_list, id=request_id)
            result_list.append(result)
        return result_list


def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    IkfastServer_pb2_grpc.add_IkfastServicer_to_server(IkFastServer(), server)
    server.add_insecure_port('[::]:50051')
    server.start()
    server.wait_for_termination()

if __name__ == '__main__':
    logging.basicConfig()
    serve()

# joint_angles = [-3.1, -1.6, 1.6, -1.6, -1.6, 0.]  # in radians
#
# # Test forward kinematics: get end effector pose from joint angles
# print("\nTesting forward kinematics:\n")
# print("Joint angles:")
# print(joint_angles)
# ee_pose = ur5_kin.forward(joint_angles)
# ee_pose = np.asarray(ee_pose).reshape(3, 4)  # 3x4 rigid transformation matrix
# print("\nEnd effector pose:")
# print(ee_pose)
# print("\n-----------------------------")
#
# # Test inverse kinematics: get joint angles from end effector pose
# print("\nTesting inverse kinematics:\n")
# ee = ee_pose.reshape(-1).tolist()
# joint_configs = ur5_kin.inverse(ee)
# n_solutions = int(len(joint_configs)/n_joints)
# print("%d solutions found:" %(n_solutions) )
# joint_configs = np.asarray(joint_configs).reshape(n_solutions,n_joints)
# for joint_config in joint_configs:
#     print(joint_config)
#
# # Check cycle-consistency of forward and inverse kinematics
# assert(np.any([np.sum(np.abs(joint_config-np.asarray(joint_angles))) < 1e-4 for joint_config in joint_configs]))
# print("\nTest passed!")